import ssl

import torchvision.datasets as tvd
import torchvision.transforms as tvt
from torch.utils.data import DataLoader

from src.config import config


def create_mnist_loader(path, train):
    if train:
        train_data = tvd.MNIST(path, download=True, train=True, transform=tvt.ToTensor())
        return DataLoader(train_data, batch_size=50, shuffle=True, num_workers=2)
    else:
        infer_data = tvd.MNIST(path, download=True, train=False, transform=tvt.ToTensor())
        return DataLoader(infer_data, batch_size=50, shuffle=False, num_workers=2)


def create_cifar10_loader(path, train):
    if train:
        train_data = tvd.CIFAR10(path, download=False, train=True, transform=tvt.Compose([
            tvt.Resize(256),
            tvt.RandomCrop(256, padding=4),
            tvt.RandomHorizontalFlip(),
            tvt.ToTensor(),
            tvt.Normalize((0.4914, 0.4822, 0.4465),(0.2023, 0.1994, 0.2010))
        ]))
        return DataLoader(train_data, batch_size=8, shuffle=True, num_workers=0)
    else:
        infer_data = tvd.CIFAR10(path, download=False, train=False, transform=tvt.Compose([
            tvt.Resize(256),
            tvt.ToTensor(),
            tvt.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
        ]))
        return DataLoader(infer_data, batch_size=8, shuffle=True, num_workers=0)


def create_loader(name, train):
    if name == "mnist":
        return create_mnist_loader(config.mnist_root, train)
    elif name == "cifar10":
        return create_cifar10_loader(config.cifar10_root, train)
    else:
        raise ValueError(f"unsupported dataset {name}.")